#include <Math/SMatrix.h>
#include <Math/SVector.h>
#include <TAxis.h>
#include <TRandom3.h>
#include <math.h>
#include <TAxis3D.h>
#include <TCanvas.h>
#include <TView3D.h>
#include <TLegend.h>
#include <TH2D.h>
#include <TH3F.h>
#include <TGraph.h>
#include <fstream>
#include <iomanip>
#include "TRestTask.h"
#include "TRestRun.h"
#include "TRestG4Metadata.h"
#include "TRestG4Event.h"
#include "TRestHitsEvent.h"
#include "TFile.h"
#include "TTree.h"

#include "toolkitOfRest.C"
TRandom3 gaus;
TRandom3 rnd;

const float XY_PRECISION = 1;//mm
const float Z_PRECISION = 1;

void X_KF3D_PyData_Gamma()
{
	
	typedef ROOT::Math::SMatrix<double,6,6>  SMatrix66;
	typedef ROOT::Math::SMatrix<double,3,6>  SMatrix36;
	typedef ROOT::Math::SMatrix<double,6,3>  SMatrix63;
	typedef ROOT::Math::SMatrix<double,3,3>  SMatrix33;
	typedef ROOT::Math::SVector<double,6>  SVector6;
	typedef ROOT::Math::SVector<double,3>  SVector3;
	typedef ROOT::Math::SVector<double,3>  SVector3;

	double MeasureError = 0.3;
	double MeasureErrorx = 0.3;
	double MeasureErrory = 0.3;
	TCanvas *c1 = new TCanvas("c1","c1",800,800);;
	//Load root file
    string mana_filename = "/lustre/pandax/wangxuan/restTest/DataEnhance/KFData/KFData_e_z-_wxenergy.5MeV_n1w_cc.root";
	string outRootName = "/lustre/pandax/wangxuan/restTest/MeGat/dataRoot/KF_e_z-_wxenergy.5MeV_n1w_cc.root";
	
    cout << "Manafile loaded : " << mana_filename << endl;

    ///////////////////////
    //Read Tree*
    TFile* inputFile = new TFile(mana_filename.c_str());
    TTree* kfdata_tree = (TTree*)inputFile->Get("KFData");
	TArrayD *vec_id = nullptr;
	TArrayD *fx = nullptr;
	TArrayD *fy = nullptr;
	TArrayD *fz = nullptr;
	TArrayD *electron_energy = nullptr;
	TArrayD *G4_e_vertex_x = nullptr;
	TArrayD *G4_e_vertex_y = nullptr;
	TArrayD *G4_e_vertex_z = nullptr;
	TArrayD *t_gamma_x = nullptr;
	TArrayD *t_gamma_y = nullptr;
	TArrayD *t_gamma_z = nullptr;
	TArrayD *totalDepositedEnergy_smear = nullptr;
	
	kfdata_tree->SetBranchAddress("id", &vec_id);
	kfdata_tree->SetBranchAddress("fx", &fx);
	kfdata_tree->SetBranchAddress("fy", &fy);
	kfdata_tree->SetBranchAddress("fz", &fz);
	kfdata_tree->SetBranchAddress("electron_energy", &electron_energy);
	kfdata_tree->SetBranchAddress("G4_e_vertex_x", &G4_e_vertex_x);
	kfdata_tree->SetBranchAddress("G4_e_vertex_y", &G4_e_vertex_y);
	kfdata_tree->SetBranchAddress("G4_e_vertex_z", &G4_e_vertex_z);
	kfdata_tree->SetBranchAddress("gamma_x", &t_gamma_x);
	kfdata_tree->SetBranchAddress("gamma_y", &t_gamma_y);
	kfdata_tree->SetBranchAddress("gamma_z", &t_gamma_z);
	kfdata_tree->SetBranchAddress("totalDepositedEnergy_smear", &totalDepositedEnergy_smear);

    int num_entries_KFData = kfdata_tree->GetEntries();
    cout << "the number of total KFData entries is : " << num_entries_KFData << endl;
    /////////////////////////////////////////////////////////////////////

	int eventID = 0;
    float filter_x = 0.;
	float filter_y = 0.;
	float filter_z = 0.;
	float Cluser_x0 = 0.;
	float Cluser_y0 = 0.;
	float Cluser_z0 = 0.;
    vector<float> mana_fx;
	vector<float> mana_fy;
	vector<float> mana_fz;
	float e_energy = 0.;
    float dz = 1;
    
	
    TFile * outRootfile = new TFile(outRootName.c_str() , "RECREATE");
	TTree * outTree  = new TTree( "KFTree" , "Kalman Filter Data" ); 
	float KFx_0;
	float KFy_0;
	float KFz_0;
	float kfvx0;
	float kfvy0;
	float kfvz0;
	float gamma_x;
	float gamma_y;
	float gamma_z;
	float gamma_energy;
	vector<float> KF_x;
	vector<float> KF_y;
	vector<float> KF_z;
	float g4_e_vertex_x;
	float g4_e_vertex_y;
	float g4_e_vertex_z;
	float vector_electron[3];
	float vertex_distance_G4_Cluser;
	float vertex_distance_G4_KF;
	// float kf_arm_x = 0;
	// float kf_arm_y = 0;
	float dz_avg = 0.;;
	outTree->Branch("id", & eventID);

	// 滤波前的点
	outTree->Branch("Cluser_x0", & Cluser_x0);
	outTree->Branch("Cluser_y0", & Cluser_y0);
	outTree->Branch("Cluser_z0", & Cluser_z0);
	outTree->Branch("Cluser_x", & mana_fx);
	outTree->Branch("Cluser_y", & mana_fy);
	outTree->Branch("Cluser_z", & mana_fz);

	// 滤波后的点
	outTree->Branch("KF_x", & KF_x);
	outTree->Branch("KF_y", & KF_y);
	outTree->Branch("KF_z", & KF_z);
	
	// 滤波后的第一个点
	outTree->Branch("KFx_0", & KFx_0);
	outTree->Branch("KFy_0", & KFy_0);
	outTree->Branch("KFz_0", & KFz_0);
	// 滤波后的初始速度
	outTree->Branch("KFvx_0", & vector_electron[0]);
	outTree->Branch("KFvy_0", & vector_electron[1]);
	outTree->Branch("KFvz_0", & vector_electron[2]);
	
	// G4中的顶点
	outTree->Branch("G4_e_vertex_x", & g4_e_vertex_x);
	outTree->Branch("G4_e_vertex_y", & g4_e_vertex_y);
	outTree->Branch("G4_e_vertex_z", & g4_e_vertex_z);
	
	outTree->Branch("gamma_x", & gamma_x);
	outTree->Branch("gamma_y", & gamma_y);
	outTree->Branch("gamma_z", & gamma_z);

	outTree->Branch("vertex_distance_G4_Cluser", &vertex_distance_G4_Cluser);
	outTree->Branch("vertex_distance_G4_KF", &vertex_distance_G4_KF);

	outTree->Branch("e_energy", & e_energy);
	outTree->Branch("gamma_energy", & gamma_energy);
	outTree->Branch("dz_avg", & dz_avg);


	int count = 0;
	float sum_arm2;
	// float rms_arm;
	int pic_count = 0;
	
	// 处理开始
	for (int i=0;i<num_entries_KFData;i++)
	{	
		mana_fx.clear();
		mana_fy.clear();
		mana_fz.clear();
		KF_x.clear();
		KF_y.clear();
		KF_z.clear();
		dz_avg = 0;
		
		//Mana data
		kfdata_tree->GetEntry(i);

		for(int j = 0; j < electron_energy->GetSize(); j++)
        {
			cout << "electron_energy: " << (*electron_energy)[j] << endl;
			e_energy = (*electron_energy)[j];
	
			cout << "vec_id: " << (*vec_id)[j] << endl;
			eventID = (*vec_id)[j];
			g4_e_vertex_x = (*G4_e_vertex_x)[j];
			g4_e_vertex_y = (*G4_e_vertex_y)[j];
			g4_e_vertex_z = (*G4_e_vertex_z)[j];
			gamma_x = (*t_gamma_x)[j];
			gamma_y = (*t_gamma_y)[j];
			gamma_z = (*t_gamma_z)[j];
			float totalDepositedEnergy = (*totalDepositedEnergy_smear)[j];

			gamma_energy = totalDepositedEnergy - e_energy;
			Cluser_x0  = (*fx)[j];
			Cluser_y0  = (*fy)[j];
			Cluser_z0  = (*fz)[j];
		}
		// eventID = i;
	
		for(int j = 0; j < fx->GetSize(); j++)
        {
			// if (j==0)
			// {
			// 	e_energy = (*electron_energy)[j];
			// 	cout << j << ", electron_energy: " << (*electron_energy)[j] << endl;
			// }
            filter_x = (*fx)[j];
            filter_y = (*fy)[j];
            filter_z = (*fz)[j];
 
			mana_fx.push_back(filter_x);
			mana_fy.push_back(filter_y);
			mana_fz.push_back(filter_z);
            
        }
	
		cout<<"EventID = :"<<eventID << endl;
	
		//速度信息产生办法
		vector<float> dLx, dLy, dLz, dL;
		vector<float> speedx_group, speedy_group, speedz_group;
		float thetarnd0 = 13600. / e_energy * sqrt(dz/23520.); //多次散射角
		cout<<"thetarnd0: "<<thetarnd0 <<", e_energy :"<< e_energy<<endl;
		for(int k=0;k<mana_fx.size()-1;k++)
		{	
			float speedx,speedy,speedz=0.;
			
			dLx.push_back(mana_fx[k+1]-mana_fx[k]);
			dLy.push_back(mana_fy[k+1]-mana_fy[k]);
			dLz.push_back(mana_fz[k+1]-mana_fz[k]);
			dL.push_back(sqrt(pow(dLx[k],2)+pow(dLy[k],2)+pow(dLz[k],2)));


			speedx=(gaus.Gaus(dLx[k],thetarnd0));
			speedy=(gaus.Gaus(dLy[k],thetarnd0));
			speedz=(gaus.Gaus(dLz[k],thetarnd0));

			speedx_group.push_back(speedx/sqrt(pow(speedx, 2)+pow(speedy, 2)+pow(speedz, 2)));
			speedy_group.push_back(speedy/sqrt(pow(speedx, 2)+pow(speedy, 2)+pow(speedz, 2)));
			speedz_group.push_back(speedz/sqrt(pow(speedx, 2)+pow(speedy, 2)+pow(speedz, 2)));
		}
		cout << "point: " << 0 << ", position: (" << mana_fx[0] << ", " << mana_fy[0] << ", " << mana_fz[0]
			<< "), vector: (" << speedx_group[0] << ", " << speedy_group[0] << ", " << speedz_group[0] << ")" <<endl;
		// if (abs(speedx_group[0]) > 0.5 || abs(speedy_group[0]) > 0.5) continue;

        float* Mana_fx = mana_fx.data();
        float* Mana_fy = mana_fy.data();
        float* Mana_fz = mana_fz.data();

		float* DLx = dLx.data();
		float* DLy = dLy.data();
		float* DLz = dLz.data();
		
		float* Speedx_group = speedx_group.data();
		float* Speedy_group = speedy_group.data();
		float* Speedz_group = speedz_group.data();

		//参数准备	
		double Speedx, Speedy, Speedz;
		int N = fx->GetSize();
		cout << "N = " << N << endl;

		//卡尔曼滤波点和平滑过程点
  		double KFx[N], KFy[N], KFz[N], KFvx[N], KFvy[N], KFvz[N], SRx[N], SRy[N], SRz[N], SRvx[N], SRvy[N], SRvz[N];
		double NUM[N];
		

		SVector3 XMscoli;
		SVector6 Expectation;//6D的预测vector
  		SVector6 *ExpectationRecord = new SVector6[N];
  		SVector6 Kalmanresult;//滤波
  		SVector6 SmoothResult;//平滑

  		SMatrix66 F;//6*6的状态矩阵F
  		SMatrix36 measureMatrix;//3*6的测量矩阵H
  		measureMatrix(0,0) = 1.0;//H的初始值
  		measureMatrix(1,1) = 1.0;
  		measureMatrix(2,2) = 1.0;
  		SMatrix66 I;//6*6的单位阵
  		I(0,0) = 1.0;
  		I(1,1) = 1.0;
  		I(2,2) = 1.0;
  		I(3,3) = 1.0;
  		I(4,4) = 1.0;
		I(5,5) = 1.0;
	
		SMatrix66 Priori;
  		SMatrix66 *PrioriRecord = new SMatrix66[N];
  		SMatrix66 Posteriori;
  		SMatrix66 *PosterioriRecord = new SMatrix66[N];
  		SMatrix63 KalmanGain;
  		SMatrix66 SmoothGain;
  		SMatrix66 NoiseThMatrix;
  		SMatrix33 NoiseMsMatrix;
  		SMatrix33 KIn;

		//参数初始化
  		Posteriori(0,0)=1.0;
  		Posteriori(1,1)=1.0;
  		Posteriori(2,2)=1.0;
  		Posteriori(3,3)=1.0;
  		Posteriori(4,4)=1.0;
  		Posteriori(5,5)=1.0;
	
		PosterioriRecord[0] = Posteriori;
	
		KFx[0] = Mana_fx[0];
  		KFy[0] = Mana_fy[0];
  		KFz[0] = Mana_fz[0];
  		KFvx[0] = Speedx_group[0];//speedx_group[0]
  		KFvy[0] = Speedy_group[0];//speedy_group[0]
  		KFvz[0] = Speedz_group[0];//speedz_group[0]	

		//卡尔曼滤波过程for循环
		for(int n2=1; n2 < N; n2++){
			Speedx = Speedx_group[n2 - 1];
			Speedy = Speedy_group[n2 - 1];
			Speedz = Speedz_group[n2 - 1];
		
			Kalmanresult[0] = KFx[n2 - 1];
    		Kalmanresult[1] = KFy[n2 - 1];
    		Kalmanresult[2] = KFz[n2 - 1];
    		Kalmanresult[3] = KFvx[n2 - 1];
    		Kalmanresult[4] = KFvy[n2 - 1];
    		Kalmanresult[5] = KFvz[n2 - 1];
		
			F(0,0) = 1.0;
        	F(0,3) = DLz[n2-1] / Speedz;//dt 
        	F(1,1) = 1.0;
        	F(1,4) = DLz[n2-1] / Speedz;
        	F(2,2) = 1.0;
        	F(2,5) = DLz[n2-1] / Speedz;
        	F(3,3) = 1.0;
        	F(4,4) = 1.0;
        	F(5,5) = 1.0;
		
    		Expectation = F * Kalmanresult;

    		ExpectationRecord[n2] = Expectation;//x(k+1)=x(k)+vx*dt

			// float thetarnd0 = 13600. / e_energy * sqrt(DLz[n2-1]/23520.);
    		NoiseThMatrix(0,0) = 0.0;
    		NoiseThMatrix(1,1) = 0.0;
    		NoiseThMatrix(2,2) = 0.0;
    		NoiseThMatrix(3,3) = 2.0 * sin(thetarnd0/2) * sqrt(0.5 * (pow(Speedy, 2) + pow(Speedz, 2)));
    		NoiseThMatrix(4,4) = 2.0 * sin(thetarnd0/2) * sqrt(0.5 * (pow(Speedx, 2) + pow(Speedz, 2)));
    		NoiseThMatrix(5,5) = 2.0 * sin(thetarnd0/2) * sqrt(0.5 * (pow(Speedx, 2) + pow(Speedy, 2)));

    		Priori = F * Posteriori*ROOT::Math::Transpose(F) + NoiseThMatrix * NoiseThMatrix;

    		PrioriRecord[n2] = Priori;//every state save in PrioriRecord 

    		NoiseMsMatrix(0,0) = MeasureErrorx;
    		NoiseMsMatrix(1,1) = MeasureErrory;
    		NoiseMsMatrix(2,2) = MeasureError;

			//cout<<"NoiseMsMatrix2: "<<NoiseMsMatrix * NoiseMsMatrix<<endl;

    		KIn = measureMatrix * Priori * ROOT::Math::Transpose(measureMatrix) + NoiseMsMatrix * NoiseMsMatrix;

    		KIn.Invert();

   	 		KalmanGain = Priori * ROOT::Math::Transpose(measureMatrix) * KIn;

    		XMscoli[0] = mana_fx[n2];
    		XMscoli[1] = mana_fy[n2];
    		XMscoli[2] = mana_fz[n2];
		
    		Kalmanresult = Expectation + KalmanGain * (XMscoli - measureMatrix * Expectation);//Kalman result
		
			Posteriori = (I - KalmanGain * measureMatrix) * Priori;
		
			PosterioriRecord[n2] = Posteriori;

			KFx[n2] = Kalmanresult[0];
			KFy[n2] = Kalmanresult[1];
			KFz[n2] = Kalmanresult[2];
			KFvx[n2] = Kalmanresult[3];
			KFvy[n2] = Kalmanresult[4];
			KFvz[n2] = Kalmanresult[5];

			NUM[n2] = n2;
		}
		//卡尔曼滤波过程结束

		//平滑过程开始
		SmoothResult[0] = KFx[N - 1];
		SmoothResult[1] = KFy[N - 1];
		SmoothResult[2] = KFz[N - 1];
		SmoothResult[3] = KFvx[N - 1];
		SmoothResult[4] = KFvy[N - 1];
		SmoothResult[5] = KFvz[N - 1];
		SRx[N - 1] = KFx[N - 1];
		SRy[N - 1] = KFy[N - 1];
		SRz[N - 1] = KFz[N - 1];
		for(int n3=N-2; n3 > -1; n3--)
		{
			Kalmanresult[0] = KFx[n3];
			Kalmanresult[1] = KFy[n3];
			Kalmanresult[2] = KFz[n3];
			Kalmanresult[3] = KFvx[n3];
			Kalmanresult[4] = KFvy[n3];
			Kalmanresult[5] = KFvz[n3];

			PrioriRecord[n3 + 1].Invert();
		
			F(0,0) = 1.0;
			F(0,3) = DLz[n3] / Speedz_group[n3];
			F(1,1) = 1.0;
			F(1,4) = DLz[n3] / Speedz_group[n3];
			F(2,2) = 1.0;
			F(2,5) = DLz[n3] / Speedz_group[n3];
			F(3,3) = 1.0;
			F(4,4) = 1.0;
			F(5,5) = 1.0;
		
			SmoothGain = PosterioriRecord[n3] * ROOT::Math::Transpose(F) * PrioriRecord[n3 + 1];

			SmoothResult = Kalmanresult + SmoothGain * (SmoothResult - ExpectationRecord[n3 + 1]);

			SRx[n3] = SmoothResult[0];
			SRy[n3] = SmoothResult[1];
			SRz[n3] = SmoothResult[2];
			SRvx[n3] = SmoothResult[3];
			SRvy[n3] = SmoothResult[4];
			SRvz[n3] = SmoothResult[5];
		
		}
		//平滑过程结束，读出数据
		float nol_SRv[3];
		for (int k=0; k < sizeof(SRx)/sizeof(SRx[0]); k++)
		{	
			nol_SRv[0] = SRvx[k];
			nol_SRv[1] = SRvy[k];
			nol_SRv[2] = SRvz[k];
			normalize_vector(nol_SRv);

			KF_x.push_back(SRx[k]);
			KF_y.push_back(SRy[k]);
			KF_z.push_back(SRz[k]);
			// cout << " SRdata: " << k << ", position: (" << SRx[k] << ", " << SRy[k] << ", " << SRz[k]
			// << "), vector: (" << nol_SRv[0] << ", " << nol_SRv[1] << ", " << nol_SRv[2] << ")" <<endl;
		}
		KFx_0 = SRx[0];
		KFy_0 = SRy[0];
		KFz_0 = SRz[0];

		nol_SRv[0] = SRvx[0];
		nol_SRv[1] = SRvy[0];
		nol_SRv[2] = SRvz[0];
		normalize_vector(nol_SRv);
		kfvx0 = nol_SRv[0];
		kfvy0 = nol_SRv[1];
		kfvz0 = nol_SRv[2];

		vector_electron[0] = kfvx0;
		vector_electron[1] = kfvy0;
		vector_electron[2] = kfvz0;
	
		
		delete[] ExpectationRecord;
		delete[] PrioriRecord;
		delete[] PosterioriRecord;

		float vector_electron_mana[3];
		vector_electron_mana[0] = Mana_fx[1] - Mana_fx[0];
		vector_electron_mana[1] = Mana_fy[1] - Mana_fy[0];
		vector_electron_mana[2] = Mana_fz[1] - Mana_fz[0];

		normalize_vector(vector_electron_mana);
		normalize_vector(vector_electron);
		
		cout << "vector_electron_0:(" << vector_electron[0] << ", " << vector_electron[1] << ", " << vector_electron[2] << ")" << endl;
	
		dz_avg = calculateAverage(dLz);
		cout << "_______________________________________________________________" << endl;
		vertex_distance_G4_Cluser = sqrt(pow((Cluser_x0 - g4_e_vertex_x),2) + pow((Cluser_y0 - g4_e_vertex_y),2) + pow((Cluser_z0 - g4_e_vertex_z),2));
		vertex_distance_G4_KF = sqrt(pow((KFx_0 - g4_e_vertex_x),2) + pow((KFy_0 - g4_e_vertex_y),2) + pow((KFz_0 - g4_e_vertex_z),2));
		outTree->Fill();

		if (pic_count < 20)
		{
			pic_count++;
		// 绘图
		// XZ

			TGraph *gr21 = new TGraph (mana_fx.size(), Mana_fx, Mana_fz);
			gr21->GetXaxis()->SetTitle("Y [mm]");
			gr21->GetYaxis()->SetTitle("Z [mm]");
			gr21->SetTitle("Project YZ");
			gr21->SetMarkerStyle(7);
			gr21->SetMarkerSize(0.5);
			gr21->SetMarkerColor(1);
			gr21->Draw("AP");

			TGraph *gr22 = new TGraph (N,SRx,SRz);
			gr22->SetMarkerStyle(7);
			gr22->SetMarkerSize(0.5);
			gr22->SetMarkerColor(2);
			gr22->Draw("P""same");

			TGraph *gr23 = new TGraph (N,SRx,SRz);
			gr23->SetLineWidth(1);
			gr23->SetLineColor(2);
			gr23->Draw("C""same");

			string xz_name = "/lustre/pandax/wangxuan/restTest/MeGat/subScript/KF_3D_e_wxenergy.5/pic/xz/mergexz_" + to_string(i) + ".png";
			c1->Print(xz_name.c_str());
			// YZ
			TGraph *gr31 = new TGraph (mana_fx.size(), Mana_fy, Mana_fz);
			gr31->GetXaxis()->SetTitle("Y [mm]");
			gr31->GetYaxis()->SetTitle("Z [mm]");
			gr31->SetTitle("Project YZ");
			gr31->SetMarkerStyle(7);
			gr31->SetMarkerSize(0.5);
			gr31->SetMarkerColor(1);
			gr31->Draw("AP");

			TGraph *gr32 = new TGraph (N,SRy,SRz);
			gr32->SetMarkerStyle(7);
			gr32->SetMarkerSize(0.5);
			gr32->SetMarkerColor(2);
			gr32->Draw("P""same");

			TGraph *gr33 = new TGraph (N,SRy,SRz);
			gr33->SetLineWidth(1);
			gr33->SetLineColor(2);
			gr33->Draw("C""same");

			string yz_name = "/lustre/pandax/wangxuan/restTest/MeGat/subScript/KF_3D_e_wxenergy.5/pic/yz/mergeyz_" + to_string(i) + ".png";
			c1->Print(yz_name.c_str());
		}

		mana_fx.clear();
		mana_fy.clear();
		mana_fz.clear();
		dLx.clear();
		dLy.clear();
		dLz.clear();
		speedx_group.clear();
		speedy_group.clear();
		speedz_group.clear();
	}
	// rms_arm = sum_arm2 / count;
	// cout << "recorded entries number: " << count << ", rms_arm_0: " << rms_arm << endl;
	outTree->Write();
    outRootfile->Close();
}
